"""
Browser Manager - Coordinate multiple browser instances.

This module manages browser instance lifecycle, session state application,
and coordinates multiple browser launches for different websites.
"""

import os
import sys
import time
import subprocess
import json
import logging
import shutil
from typing import Dict, List, Optional, Any, Tuple
from pathlib import Path
from dataclasses import dataclass

from models import WebsiteProfile, UserSession, BrowserConfig
from auth_manager import AuthenticationManager


@dataclass
class BrowserInstance:
    """Represents a running browser instance."""
    process_id: int
    website_id: str
    username: str
    profile_dir: str
    target_url: str
    process: subprocess.Popen
    
    def is_running(self) -> bool:
        """Check if browser process is still running."""
        return self.process.poll() is None
    
    def terminate(self):
        """Terminate the browser process."""
        try:
            self.process.terminate()
        except Exception:
            pass


class BrowserManager:
    """
    Manages browser instance lifecycle and coordination.
    
    Handles launching multiple browser instances, applying session states,
    managing browser isolation, and coordinating resource usage.
    """
    
    def __init__(self, profiles_base_dir: str = "browser_profiles"):
        """
        Initialize the Browser Manager.
        
        Args:
            profiles_base_dir: Base directory for browser profiles
        """
        self.profiles_base_dir = Path(profiles_base_dir)
        self.profiles_base_dir.mkdir(exist_ok=True)
        
        self.logger = self._setup_logging()
        self.auth_manager = AuthenticationManager()
        
        # Track running browser instances
        self.running_instances: Dict[str, BrowserInstance] = {}
        
        # Find Chrome executable
        self.chrome_path = self.auth_manager.chrome_path
        if not self.chrome_path:
            raise RuntimeError("Chrome browser not found")
        
        self.logger.info("BrowserManager initialized")
    
    def _setup_logging(self) -> logging.Logger:
        """Setup logging configuration."""
        logger = logging.getLogger("BrowserManager")
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def _prepare_profile_directory(self, username: str, website_id: str, 
                                 session: Optional[UserSession] = None,
                                 isolated: bool = False) -> str:
        """
        Prepare browser profile directory with session state.
        
        Args:
            username: User identifier
            website_id: Website identifier
            session: User session with auth state
            isolated: Whether to create isolated profile
            
        Returns:
            Path to prepared profile directory
        """
        try:
            if isolated:
                # Create isolated profile for this specific instance
                timestamp = int(time.time())
                profile_dir = self.profiles_base_dir / username / f"{website_id}_isolated_{timestamp}"
            else:
                # Use shared profile for all instances of this user/website
                profile_dir = self.profiles_base_dir / username / f"{website_id}_shared"
            
            profile_dir.mkdir(parents=True, exist_ok=True)
            
            # Initialize marker file
            init_marker = profile_dir / ".initialized"
            need_init = not init_marker.exists()
            
            if need_init and session:
                self.logger.info(f"Initializing profile directory: {profile_dir}")
                
                # Apply session state using Playwright
                if session.auth_state:
                    success = self._apply_session_state(str(profile_dir), session)
                    if success:
                        self.logger.info("Session state applied successfully")
                    else:
                        self.logger.warning("Failed to apply session state")
                
                # Download and apply extensions data
                try:
                    extensions_success = self.auth_manager.download_extensions_backup(
                        username, str(profile_dir)
                    )
                    if extensions_success:
                        self.logger.info("Extensions data applied successfully")
                except Exception as e:
                    self.logger.warning(f"Extensions data not available: {e}")
                
                # Create initialization marker
                with open(init_marker, 'w') as f:
                    f.write(f"Initialized at {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
                    f.write(f"Username: {username}\n")
                    f.write(f"Website: {website_id}\n")
                    f.write(f"Session ID: {session.session_id if session else 'None'}\n")
            
            return str(profile_dir)
            
        except Exception as e:
            self.logger.error(f"Failed to prepare profile directory: {e}")
            raise
    
    def _apply_session_state(self, profile_dir: str, session: UserSession) -> bool:
        """
        Apply session state to browser profile using Playwright.
        
        Args:
            profile_dir: Browser profile directory
            session: User session with auth state
            
        Returns:
            True if successful, False otherwise
        """
        try:
            from playwright.sync_api import sync_playwright
            
            with sync_playwright() as p:
                # Launch persistent context with the profile directory
                context = p.chromium.launch_persistent_context(
                    user_data_dir=profile_dir,
                    headless=True,
                    channel="chrome",
                    args=['--disable-web-security', '--no-first-run', '--no-default-browser-check']
                )
                
                # Create page and apply auth state
                page = context.new_page()
                
                # Apply cookies
                cookies = session.auth_state.get('cookies', [])
                if cookies:
                    try:
                        valid_cookies = []
                        for cookie in cookies:
                            if all(key in cookie for key in ['name', 'value', 'domain']):
                                valid_cookies.append(cookie)
                        
                        if valid_cookies:
                            context.add_cookies(valid_cookies)
                            self.logger.debug(f"Applied {len(valid_cookies)} cookies")
                    except Exception as e:
                        self.logger.warning(f"Failed to apply cookies: {e}")
                
                # Apply localStorage and sessionStorage
                origins = session.auth_state.get('origins', [])
                for origin_data in origins:
                    origin = origin_data.get('origin')
                    local_storage = origin_data.get('localStorage', [])
                    session_storage = origin_data.get('sessionStorage', [])
                    
                    if origin and (local_storage or session_storage):
                        try:
                            page.goto(origin, wait_until='domcontentloaded', timeout=10000)
                            
                            # Set localStorage
                            for item in local_storage:
                                if 'name' in item and 'value' in item:
                                    page.evaluate(f"localStorage.setItem('{item['name']}', '{item['value']}')")
                            
                            # Set sessionStorage
                            for item in session_storage:
                                if 'name' in item and 'value' in item:
                                    page.evaluate(f"sessionStorage.setItem('{item['name']}', '{item['value']}')")
                            
                            self.logger.debug(f"Applied storage data for {origin}")
                            
                        except Exception as e:
                            self.logger.warning(f"Failed to apply storage for {origin}: {e}")
                
                context.close()
                return True
                
        except ImportError:
            self.logger.error("Playwright not available for session state application")
            return False
        except Exception as e:
            self.logger.error(f"Failed to apply session state: {e}")
            return False
    
    def launch_browser(self, profile: WebsiteProfile, session: Optional[UserSession] = None,
                      username: Optional[str] = None, isolated: bool = False,
                      count: int = 1) -> List[BrowserInstance]:
        """
        Launch browser instance(s) for specified website.
        
        Args:
            profile: Website profile configuration
            session: User session with auth state
            username: User identifier (required if session not provided)
            isolated: Whether to create isolated browser instances
            count: Number of browser instances to launch
            
        Returns:
            List of launched browser instances
        """
        try:
            if not session and not username:
                raise ValueError("Either session or username must be provided")
            
            user_id = session.username if session else username
            self.logger.info(f"Launching {count} browser instance(s) for {profile.display_name}")
            
            launched_instances = []
            
            for i in range(count):
                instance_id = f"{user_id}_{profile.website_id}_{i+1}_{int(time.time())}"
                
                try:
                    # Prepare profile directory
                    profile_dir = self._prepare_profile_directory(
                        user_id, profile.website_id, session, isolated
                    )
                    
                    # Build Chrome arguments
                    chrome_args = self._build_chrome_args(profile, profile_dir)
                    
                    # Launch Chrome process
                    process = subprocess.Popen(
                        chrome_args,
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                        stdin=subprocess.DEVNULL,
                        creationflags=subprocess.CREATE_NEW_PROCESS_GROUP if sys.platform == 'win32' else 0
                    )
                    
                    # Create browser instance object
                    browser_instance = BrowserInstance(
                        process_id=process.pid,
                        website_id=profile.website_id,
                        username=user_id,
                        profile_dir=profile_dir,
                        target_url=profile.target_url,
                        process=process
                    )
                    
                    # Track instance
                    self.running_instances[instance_id] = browser_instance
                    launched_instances.append(browser_instance)
                    
                    self.logger.info(f"Launched browser instance {i+1}/{count} (PID: {process.pid})")
                    
                    # Add delay between launches to avoid resource conflicts
                    if i < count - 1:
                        time.sleep(2)
                        
                except Exception as e:
                    self.logger.error(f"Failed to launch browser instance {i+1}: {e}")
                    continue
            
            self.logger.info(f"Successfully launched {len(launched_instances)}/{count} browser instances")
            return launched_instances
            
        except Exception as e:
            self.logger.error(f"Failed to launch browsers: {e}")
            return []
    
    def _build_chrome_args(self, profile: WebsiteProfile, profile_dir: str) -> List[str]:
        """
        Build Chrome command line arguments.
        
        Args:
            profile: Website profile configuration
            profile_dir: Browser profile directory
            
        Returns:
            List of Chrome command line arguments
        """
        chrome_args = [
            self.chrome_path,
            f"--user-data-dir={profile_dir}",
            profile.target_url
        ]
        
        # Add browser configuration arguments
        browser_config = profile.browser_config
        
        # Viewport settings
        if browser_config.viewport:
            width = browser_config.viewport.get("width", 1920)
            height = browser_config.viewport.get("height", 1080)
            chrome_args.append(f"--window-size={width},{height}")
        
        # Start maximized (override window size if needed)
        chrome_args.append("--start-maximized")
        
        # User agent
        if browser_config.user_agent:
            chrome_args.append(f"--user-agent={browser_config.user_agent}")
        
        # Additional arguments
        if browser_config.extra_args:
            chrome_args.extend(browser_config.extra_args)
        
        # Default Chrome arguments for stability
        default_args = [
            "--no-first-run",
            "--no-default-browser-check"
        ]
        chrome_args.extend(default_args)
        
        return chrome_args
    
    def launch_multi_website_browsers(self, sessions: List[UserSession], 
                                    profiles: Dict[str, WebsiteProfile],
                                    count_per_website: int = 1,
                                    isolated: bool = False) -> Dict[str, List[BrowserInstance]]:
        """
        Launch browsers for multiple websites simultaneously.
        
        Args:
            sessions: List of user sessions
            profiles: Dictionary of website profiles
            count_per_website: Number of browser instances per website
            isolated: Whether to create isolated instances
            
        Returns:
            Dictionary mapping website IDs to browser instances
        """
        try:
            launched_browsers = {}
            
            for session in sessions:
                website_id = session.website_id
                
                if website_id not in profiles:
                    self.logger.warning(f"Profile not found for website: {website_id}")
                    continue
                
                profile = profiles[website_id]
                
                # Launch browsers for this website
                instances = self.launch_browser(
                    profile=profile,
                    session=session,
                    isolated=isolated,
                    count=count_per_website
                )
                
                if instances:
                    launched_browsers[website_id] = instances
                    self.logger.info(f"Launched {len(instances)} instances for {profile.display_name}")
                else:
                    self.logger.error(f"Failed to launch instances for {profile.display_name}")
            
            total_instances = sum(len(instances) for instances in launched_browsers.values())
            self.logger.info(f"Total browsers launched: {total_instances}")
            
            return launched_browsers
            
        except Exception as e:
            self.logger.error(f"Failed to launch multi-website browsers: {e}")
            return {}
    
    def get_running_instances(self) -> Dict[str, BrowserInstance]:
        """
        Get currently running browser instances.
        
        Returns:
            Dictionary of running browser instances
        """
        # Clean up terminated instances
        terminated_instances = []
        for instance_id, instance in self.running_instances.items():
            if not instance.is_running():
                terminated_instances.append(instance_id)
        
        for instance_id in terminated_instances:
            del self.running_instances[instance_id]
        
        return self.running_instances.copy()
    
    def terminate_instances(self, instance_ids: Optional[List[str]] = None) -> int:
        """
        Terminate browser instances.
        
        Args:
            instance_ids: Specific instance IDs to terminate (None for all)
            
        Returns:
            Number of instances terminated
        """
        try:
            if instance_ids is None:
                instance_ids = list(self.running_instances.keys())
            
            terminated_count = 0
            
            for instance_id in instance_ids:
                if instance_id in self.running_instances:
                    instance = self.running_instances[instance_id]
                    
                    try:
                        instance.terminate()
                        terminated_count += 1
                        self.logger.info(f"Terminated browser instance: {instance_id}")
                    except Exception as e:
                        self.logger.error(f"Failed to terminate instance {instance_id}: {e}")
                    
                    # Remove from tracking
                    del self.running_instances[instance_id]
            
            return terminated_count
            
        except Exception as e:
            self.logger.error(f"Failed to terminate instances: {e}")
            return 0
    
    def terminate_user_instances(self, username: str) -> int:
        """
        Terminate all browser instances for a specific user.
        
        Args:
            username: User identifier
            
        Returns:
            Number of instances terminated
        """
        user_instance_ids = [
            instance_id for instance_id, instance in self.running_instances.items()
            if instance.username == username
        ]
        
        return self.terminate_instances(user_instance_ids)
    
    def terminate_website_instances(self, website_id: str) -> int:
        """
        Terminate all browser instances for a specific website.
        
        Args:
            website_id: Website identifier
            
        Returns:
            Number of instances terminated
        """
        website_instance_ids = [
            instance_id for instance_id, instance in self.running_instances.items()
            if instance.website_id == website_id
        ]
        
        return self.terminate_instances(website_instance_ids)
    
    def get_instance_statistics(self) -> Dict[str, Any]:
        """
        Get statistics about running browser instances.
        
        Returns:
            Dictionary with instance statistics
        """
        running_instances = self.get_running_instances()
        
        total_instances = len(running_instances)
        websites = set(instance.website_id for instance in running_instances.values())
        users = set(instance.username for instance in running_instances.values())
        
        # Count by website
        website_counts = {}
        user_counts = {}
        
        for instance in running_instances.values():
            website_counts[instance.website_id] = website_counts.get(instance.website_id, 0) + 1
            user_counts[instance.username] = user_counts.get(instance.username, 0) + 1
        
        return {
            "total_instances": total_instances,
            "active_websites": len(websites),
            "active_users": len(users),
            "instances_by_website": website_counts,
            "instances_by_user": user_counts,
            "timestamp": time.strftime('%Y-%m-%d %H:%M:%S')
        }
    
    def cleanup_profile_directories(self, max_age_days: int = 7) -> int:
        """
        Clean up old browser profile directories.
        
        Args:
            max_age_days: Maximum age in days for profile directories
            
        Returns:
            Number of directories cleaned up
        """
        try:
            cleanup_count = 0
            current_time = time.time()
            max_age_seconds = max_age_days * 24 * 3600
            
            for user_dir in self.profiles_base_dir.iterdir():
                if not user_dir.is_dir():
                    continue
                
                for profile_dir in user_dir.iterdir():
                    if not profile_dir.is_dir():
                        continue
                    
                    # Check if profile is old enough and not currently in use
                    dir_age = current_time - profile_dir.stat().st_mtime
                    
                    if dir_age > max_age_seconds:
                        # Check if any running instance uses this profile
                        profile_path = str(profile_dir)
                        in_use = any(
                            instance.profile_dir == profile_path and instance.is_running()
                            for instance in self.running_instances.values()
                        )
                        
                        if not in_use:
                            try:
                                shutil.rmtree(profile_dir)
                                cleanup_count += 1
                                self.logger.info(f"Cleaned up old profile directory: {profile_dir}")
                            except Exception as e:
                                self.logger.warning(f"Failed to clean up {profile_dir}: {e}")
            
            self.logger.info(f"Cleaned up {cleanup_count} old profile directories")
            return cleanup_count
            
        except Exception as e:
            self.logger.error(f"Failed to cleanup profile directories: {e}")
            return 0